"""
猫狗识别项目 - Gradio网页应用
提供图片上传和预测功能的交互式界面
"""

import os
import numpy as np
from PIL import Image
import joblib
import faiss
import gradio as gr
import warnings
warnings.filterwarnings('ignore')

# ==================== 全局配置 ====================
IMAGE_SIZE = (224, 224)  # 图片尺寸
MODEL_PATH = './saved_models/ensemble_model.joblib'
SCALER_PATH = './saved_models/scaler.joblib'
PCA_PATH = './saved_models/pca.joblib'
INDEX_PATH = './saved_indexes/faiss_index.faiss'
LABELS_PATH = './saved_indexes/train_labels.npy'

# ==================== 加载模型和索引 ====================
print("=" * 50)
print("正在加载模型和索引...")
print("=" * 50)

# 全局变量
ensemble_model = None
scaler = None
pca = None
faiss_index = None
train_labels = None

def load_models():
    """
    加载训练好的模型、标准化器、PCA、Faiss索引和训练标签
    
    返回:
        success: 是否成功加载
        message: 加载信息
    """
    global ensemble_model, scaler, pca, faiss_index, train_labels
    
    try:
        # 检查文件是否存在
        if not os.path.exists(MODEL_PATH):
            return False, f"错误: 模型文件不存在 - {MODEL_PATH}\n请先运行 train_ensemble.py 训练模型!"
        
        if not os.path.exists(SCALER_PATH):
            return False, f"错误: 标准化器文件不存在 - {SCALER_PATH}\n请先运行 train_ensemble.py 训练模型!"
        
        if not os.path.exists(PCA_PATH):
            return False, f"错误: PCA文件不存在 - {PCA_PATH}\n请先运行 train_ensemble.py 训练模型!"
        
        if not os.path.exists(INDEX_PATH):
            return False, f"错误: Faiss索引文件不存在 - {INDEX_PATH}\n请先运行 train_ensemble.py 训练模型!"
        
        if not os.path.exists(LABELS_PATH):
            return False, f"错误: 训练标签文件不存在 - {LABELS_PATH}\n请先运行 train_ensemble.py 训练模型!"
        
        # 加载集成学习模型
        print("加载集成学习模型...")
        ensemble_model = joblib.load(MODEL_PATH)
        print(f"✓ 模型加载成功: {MODEL_PATH}")
        
        # 加载标准化器
        print("加载标准化器...")
        scaler = joblib.load(SCALER_PATH)
        print(f"✓ 标准化器加载成功: {SCALER_PATH}")
        
        # 加载PCA模型
        print("加载PCA模型...")
        pca = joblib.load(PCA_PATH)
        print(f"✓ PCA模型加载成功: {PCA_PATH}")
        
        # 加载Faiss索引
        print("加载Faiss索引...")
        faiss_index = faiss.read_index(INDEX_PATH)
        print(f"✓ Faiss索引加载成功: {INDEX_PATH}")
        print(f"  索引中的向量数量: {faiss_index.ntotal}")
        
        # 加载训练标签
        print("加载训练标签...")
        train_labels = np.load(LABELS_PATH)
        print(f"✓ 训练标签加载成功: {LABELS_PATH}")
        print(f"  标签数量: {len(train_labels)}")
        
        print("=" * 50)
        print("所有模型和索引加载完成!")
        print("=" * 50)
        
        return True, "模型加载成功!"
        
    except Exception as e:
        error_msg = f"加载模型时发生错误: {str(e)}"
        print(error_msg)
        return False, error_msg

# 初始化时加载模型
load_success, load_message = load_models()
if not load_success:
    print("\n" + "!" * 50)
    print("警告: 模型加载失败!")
    print(load_message)
    print("!" * 50)

# ==================== 图片预处理函数 ====================
def preprocess_image(image):
    """
    对上传的图片进行预处理
    
    参数:
        image: PIL Image对象或numpy数组
    
    返回:
        processed_image: 预处理后的特征向量
    """
    try:
        # 如果是numpy数组,转换为PIL Image
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image.astype('uint8'))
        
        # 转换为RGB模式
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        # 调整图片大小为224x224像素
        image = image.resize(IMAGE_SIZE)
        
        # 转换为numpy数组并归一化到[0, 1]
        img_array = np.array(image) / 255.0
        
        # 展平为一维向量
        img_flat = img_array.flatten()
        
        # 转换为2D数组(模型需要2D输入)
        img_2d = img_flat.reshape(1, -1)
        
        return img_2d
        
    except Exception as e:
        raise Exception(f"图片预处理失败: {str(e)}")

# ==================== 预测函数 ====================
def predict_image(image):
    """
    对上传的图片进行预测
    
    参数:
        image: 上传的图片
    
    返回:
        result: 预测结果文本
    """
    # 检查模型是否加载
    if ensemble_model is None or scaler is None or pca is None:
        return "❌ 错误: 模型未加载!\n\n请先运行 train_ensemble.py 训练模型,然后重启此应用。"
    
    # 检查图片是否上传
    if image is None:
        return "⚠️ 请先上传一张图片!"
    
    try:
        # 1. 图片预处理
        print("\n" + "-" * 50)
        print("开始预测...")
        print("-" * 50)
        print("步骤1: 预处理图片...")
        processed_image = preprocess_image(image)
        print(f"✓ 预处理完成,特征维度: {processed_image.shape}")
        
        # 2. 特征标准化
        print("步骤2: 特征标准化...")
        scaled_image = scaler.transform(processed_image)
        print("✓ 标准化完成")
        
        # 3. PCA降维
        print("步骤3: PCA降维...")
        pca_image = pca.transform(scaled_image)
        print(f"✓ 降维完成,特征维度: {pca_image.shape[1]}")
        
        # 4. 使用集成模型进行预测
        print("步骤4: 集成模型预测...")
        prediction = ensemble_model.predict(pca_image)[0]
        
        # 如果模型支持概率预测,获取预测概率
        if hasattr(ensemble_model, 'predict_proba'):
            probabilities = ensemble_model.predict_proba(pca_image)[0]
            cat_prob = probabilities[0] * 100
            dog_prob = probabilities[1] * 100
            print(f"✓ 预测完成 - 猫: {cat_prob:.2f}%, 狗: {dog_prob:.2f}%")
        else:
            cat_prob = 100 if prediction == 0 else 0
            dog_prob = 100 if prediction == 1 else 0
            print(f"✓ 预测完成 - 类别: {prediction}")
        
        # 5. 使用Faiss索引查找最相似的图片
        print("步骤5: Faiss相似度检索...")
        if faiss_index is not None and train_labels is not None:
            # 转换为float32(Faiss要求)
            query_vector = scaled_image.astype(np.float32)
            
            # 搜索最近的5个邻居
            k = 5
            distances, indices = faiss_index.search(query_vector, k)
            
            # 统计最近邻的标签
            neighbor_labels = train_labels[indices[0]]
            cat_neighbors = np.sum(neighbor_labels == 0)
            dog_neighbors = np.sum(neighbor_labels == 1)
            
            print(f"✓ 检索完成 - 最近{k}个邻居中: {cat_neighbors}个猫, {dog_neighbors}个狗")
        else:
            cat_neighbors = 0
            dog_neighbors = 0
        
        # 6. 生成预测结果
        print("-" * 50)
        print("预测完成!")
        print("-" * 50)
        
        # 确定预测类别
        if prediction == 0:
            result_emoji = "🐱"
            result_text = "猫"
            confidence = cat_prob
        else:
            result_emoji = "🐶"
            result_text = "狗"
            confidence = dog_prob
        
        # 构建结果字符串
        result = f"""
# 🎯 预测结果

## {result_emoji} 这是一只 **{result_text}**!

---

### 📊 预测详情

**置信度分布:**
- 🐱 猫: {cat_prob:.2f}%
- 🐶 狗: {dog_prob:.2f}%

**预测置信度:** {confidence:.2f}%

---

### 🔍 相似度分析 (Faiss检索)

在训练集中找到的最相似的5张图片:
- 🐱 猫: {cat_neighbors} 张
- 🐶 狗: {dog_neighbors} 张

---

### ✅ 预测完成
模型: 集成学习 (随机森林 + 逻辑回归)
特征: PCA降维 (500维)
"""
        
        return result
        
    except Exception as e:
        error_result = f"""
# ❌ 预测失败

发生错误: {str(e)}

请检查:
1. 图片格式是否正确 (支持 JPG, PNG 等常见格式)
2. 图片是否损坏
3. 模型文件是否完整

如果问题持续,请重新训练模型。
"""
        print(f"\n错误: {str(e)}")
        return error_result

# ==================== Gradio界面 ====================
def create_interface():
    """
    创建Gradio网页界面
    """
    # 自定义CSS样式
    custom_css = """
    .gradio-container {
        font-family: 'Arial', sans-serif;
    }
    .output-markdown {
        font-size: 16px;
    }
    """
    
    # 创建界面
    with gr.Blocks(css=custom_css, title="猫狗识别系统") as demo:
        # 标题
        gr.Markdown("""
        # 🐱🐶 猫狗识别系统
        
        基于集成学习和Faiss索引的智能图像分类系统
        
        ---
        
        ### 📝 使用说明
        1. 点击下方区域上传一张猫或狗的图片
        2. 系统会自动进行预测
        3. 查看预测结果和置信度
        
        ### 🔧 技术栈
        - **机器学习**: 随机森林 + 逻辑回归 (集成学习)
        - **特征降维**: PCA (主成分分析)
        - **相似度检索**: Faiss (Facebook AI Similarity Search)
        - **Web框架**: Gradio
        """)
        
        # 主要内容区域
        with gr.Row():
            # 左侧: 图片上传
            with gr.Column(scale=1):
                gr.Markdown("### 📤 上传图片")
                image_input = gr.Image(
                    label="上传猫或狗的图片",
                    type="pil",
                    height=400
                )
                
                predict_button = gr.Button(
                    "🚀 开始预测",
                    variant="primary",
                    size="lg"
                )
            
            # 右侧: 预测结果
            with gr.Column(scale=1):
                gr.Markdown("### 📊 预测结果")
                output_text = gr.Markdown(
                    value="等待上传图片...",
                    label="预测结果"
                )
        
        # 示例图片(如果有的话)
        gr.Markdown("""
        ---
        ### 💡 提示
        - 支持的图片格式: JPG, PNG, BMP, WEBP 等
        - 图片会自动调整为 224x224 像素
        - 建议上传清晰的猫或狗的正面照片以获得最佳效果
        """)
        
        # 绑定预测函数
        predict_button.click(
            fn=predict_image,
            inputs=image_input,
            outputs=output_text
        )
        
        # 也可以在上传图片时自动预测
        image_input.change(
            fn=predict_image,
            inputs=image_input,
            outputs=output_text
        )
    
    return demo

# ==================== 主程序 ====================
if __name__ == "__main__":
    print("\n" + "=" * 50)
    print("启动Gradio网页应用...")
    print("=" * 50)
    
    # 检查模型是否加载成功
    if not load_success:
        print("\n" + "!" * 50)
        print("警告: 模型未成功加载!")
        print("应用将启动,但预测功能可能无法使用。")
        print("请先运行 train_ensemble.py 训练模型。")
        print("!" * 50 + "\n")
    
    # 创建并启动界面
    demo = create_interface()
    
    # 启动服务器
    # share=True 会创建公共链接(可选)
    # server_name="0.0.0.0" 允许局域网访问
    demo.launch(
        server_name="127.0.0.1",  # 本地访问
        server_port=7860,          # 端口号
        share=False,               # 不创建公共链接
        show_error=True            # 显示错误信息
    )
